# SPDX-License-Identifier: Apache-2.0
# (C) Copyright IBM Corp. 2024.
# Licensed under the Apache License, Version 2.0 (the “License”);
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#  http://www.apache.org/licenses/LICENSE-2.0
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an “AS IS” BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
################################################################################

from typing import Tuple

import pyarrow as pa
from data_processing.test_support.transform.table_transform_test import (
    AbstractTableTransformTest,
)
from dpk_malware.transform import MalwareTransform


table = pa.Table.from_pydict(
    {
        "document_id": ["ID_1", "ID_2"],
        "contents": ["INNOCENT", "X5O!P%@AP[4\\PZX54(P^)7CC)7}$EICAR-STANDARD-ANTIVIRUS-TEST-FILE!$H+H*"],
    }
)
expected_table = pa.Table.from_pydict(
    {
        "document_id": ["ID_1", "ID_2"],
        "contents": ["INNOCENT", "X5O!P%@AP[4\\PZX54(P^)7CC)7}$EICAR-STANDARD-ANTIVIRUS-TEST-FILE!$H+H*"],
        "virus_detection": [None, "Win.Test.EICAR_HDB-1"],
    }
)
expected_metadata_list = [{"clean": 1, "infected": 1}, {}]  # transform() result  # flush() result


class TestMalwareTransform(AbstractTableTransformTest):
    """
    Extends the super-class to define the test data for the tests defined there.
    The name of this class MUST begin with the word Test so that pytest recognizes it as a test class.
    """

    def get_test_transform_fixtures(self) -> list[Tuple]:
        fixtures = [
            (
                MalwareTransform({"malware_input_column": "contents", "malware_output_column": "virus_detection"}),
                [table],
                [expected_table],
                expected_metadata_list,
            ),
        ]
        return fixtures
